import os
import json
import os.path as osp
from copy import deepcopy

_labelme_format = {
    "imagePath": "",
    "imageHeight":0,
    "imageWidth":0,
    "imageData":None,
    "lineColor":[
        0,
        255,
        0,
        128
    ],
    "fillColor":[
        255,
        0,
        0,
        128
    ],
    "flags":{},
    "shapes":[
        # {
        #     "line_color":None,
        #     "fill_color":None,
        #     "label":"",
        #     "shape_type":"polygon",
        #     "flags":{},
        #     "points":[],
        # }   
    ]
}

def polygon2points(polygon):
    points = []
    for i in range(0, len(polygon), 2):
        points.append(polygon[i:i+2])
    return points

def coco2labelme(annotations):
    images = annotations['images']
    image_infos = {}
    for image in images:
        labelme_format = deepcopy(_labelme_format)
        file_name = osp.basename(image['file_name'])
        height = image['height']
        width = image['width']
        labelme_format['imagePath'] = file_name
        labelme_format['imageHeight'] = height
        labelme_format['imageWidth'] = width
        image_infos[image['id']] = labelme_format
    categories = {cat['id']: cat['name'] for cat in annotations['categories']}
    segmentations = annotations['annotations']
    for segmentation in segmentations:
        image_id = segmentation['image_id']
        category_id = segmentation['category_id']
        label = categories[category_id]
        polygon = segmentation['segmentation'][0]
        points = polygon2points(polygon)
        shape = dict(line_color=None, fill_color=None, label=label, shape_type='polygon', flags={}, points=points)
        image_infos[image_id]['shapes'].append(shape)
    return image_infos

def loadjson(jspath):
    with open(jspath) as fp:
        return json.load(fp)

def makedirs(dstpath, cover=False):
    if osp.exists(dstpath):
        if cover:
            import shutil
            shutil.rmtree(dstpath)
            os.makedirs(dstpath)    
    else:
        os.makedirs(dstpath)

def savejson(jsdict, jspath):
    with open(jspath, 'w') as fp:
        json.dump(jsdict, fp=fp)

def coco2labelme_main(annpath, savepath):

    makedirs(savepath, cover=False)

    annotations = loadjson(annpath)
    image_infos = coco2labelme(annotations)
    for idx, img_info in image_infos.items():
        file_name = img_info['imagePath']
        print('[%d/%d] %s' % (idx + 1, len(image_infos), file_name))
        dstpath = osp.join(savepath, osp.splitext(file_name)[0]+'.json')
        savejson(img_info, dstpath)


if __name__ == '__main__':
    annpath = r'G:\data\scls_data\instance.json'
    savepath = r'G:\data\scls_data\instance'
    coco2labelme_main(annpath, savepath)
